Architecture Proposal

#2

Implementation Details

This PR implements SiameseNorm + Depth-Attention on top of the OLMo 3 decoder architecture and includes the reproducible Ascend/MindSpeed training pipeline used for the current model artifacts.

Modified components

• Added separate base and siamese_depth OLMo 3 model variants.
• Added forward-local dual residual-stream state for SiameseNorm.
• Added sparse cross-layer Depth-Attention before ordinary token attention.
• Preserved OLMo 3 Q/K normalization, RoPE, MLP, tokenizer, causal masking, and the [SWA, SWA, SWA, Full] attention pattern.
• Added Stage 1 pretraining, Stage 2 mid-training, Stage 3 long-context training, and Stage 4 Think/Instruct SFT configuration and launch paths.
• Added checkpoint save/resume/transition validation, native inference, 65K KV-cache validation, PPL evaluation, and OLMES objective-evaluation entry points.
• Added Transformers remote-code inference files under archs/SiameseNorm-DepthAttention.
• Added the reproducible Megatron/MindSpeed pipeline under reproduce/Megatron-LM.

Key implementation details

• SiameseNorm maintains two forward-local residual streams while sharing each attention and MLP computation. For layer l, the Post-Norm-like stream receives the depth-scaled update Δ / sqrt(2l), while the Pre-Norm-like stream receives the full update Δ.
• Depth-Attention is applied after Q/K normalization and RoPE, but before SWA or Full Attention.
• The current query selects among the current layer and sparse earlier-layer K/V states at the same token position. The default source stride is 8.
• The mixed value is passed into the original token-attention operation and written to the inference KV cache.
• Depth-Attention adds no learned parameters or separate persistent cache.
• Stage 3 and Stage 4 apply YaRN only to Full-Attention layers, supporting contexts up to 65,536 tokens.
• SWA layers retain the original RoPE and a 4,096-token attention window.
• Base and modified models use separate, strictly validated checkpoint keyspaces.
• Checkpoint save, resume, topology migration, and Stage 1 → Stage 2 → Stage 3 → Stage 4 transitions are included in the pipeline.

Compatibility considerations

• The published 1B artifacts use standard Transformers remote-code format and were release-validated with BF16 SDPA.
• The native training path targets Megatron/MindSpeed on Ascend NPUs.
• Third-party revisions are pinned, while deployment-specific paths, cluster configuration, and credentials remain outside version control.
• Model construction is available for the repository's OLMo 3 1B, 3B, and 7B presets.
• End-to-end four-stage training and checkpoint publication have been completed for the modified 1B model.
• A matched OLMo 3 1B baseline has also completed, enabling a controlled preliminary comparison.
• The modified 3B experiment is still training.
• The 7B experimental plan is currently suspended.

Current scope and limitations

• This remains a draft implementation PR, not a completed multi-scale architectural-effectiveness claim.
• The modified four-stage checkpoints of OLMo 3 1B and its matched baseline have been published, and its downstream evaluation are complete.
• Training curves indicate faster early-stage convergence for SiameseNorm + Depth-Attention relative to the matched baseline.
• However, the final 1B downstream difference is small: the modified model reaches a 36.9 eight-task macro average, compared with 36.8 for the matched OLMo 3 baseline.
• The 1B results are therefore mixed and do not yet establish a substantial final quality improvement.
• The modified 3B experiment remains in progress.
• The 7B experimental plan is suspended and is not part of the currently active training schedule.

Experimental Validation

Research Question 1 — Does the implementation preserve OLMo 3 training, checkpoint, and long-context inference semantics?

Hypothesis: SiameseNorm and Depth-Attention can be integrated without breaking the four-stage OLMo 3 pipeline, checkpoint transitions, SWA/Full-Attention behavior, or 65K inference.

Results & Analysis:

The modified 1B model completed and published the following artifacts:

Checkpoint save and resume, stage transitions, long-context attention, SWA/Full-Attention RoPE separation, and native inference were exercised through the complete 1B pipeline.

Findings: The current evidence supports implementation correctness and end-to-end operability at the 1B scale.

Research Question 2 — Does SiameseNorm + Depth-Attention improve model quality?

Hypothesis: At matched model size, data order, token count, optimizer, learning-rate schedule, batch size, random seed, and evaluation protocol, the modified architecture will improve convergence and downstream evaluation without introducing instability.

Results & Analysis:

The matched 1B training curves show faster loss reduction for SiameseNorm + Depth-Attention during the earlier portion of training. This indicates a potential optimization or sample-efficiency benefit.

However, the final objective downstream evaluation is nearly tied. The modified model reaches a macro average of 36.9, compared with 36.8 for the matched OLMo 3 1B baseline, corresponding to an aggregate difference of approximately +0.1 points.

The modified model outperforms the matched OLMo 3 baseline on four tasks:

• IFEval: +6.5
• TruthfulQA: +3.2
• MATH: +2.0
• BBH: +1.6

It underperforms on the remaining four tasks:

• DROP: -4.9
• MMLU: -4.2
• GSM8K: -2.0
• PopQA: -1.7

The published OLMo 2 1B pure-SFT results are retained only as an external reference. They are not a matched architectural control because the model family and training recipe differ.

Findings: SiameseNorm + Depth-Attention shows faster early-stage convergence and meaningful gains on selected tasks, particularly IFEval. Nevertheless, the final 1B aggregate result is only marginally above the matched baseline. The current evidence supports an early-convergence benefit but does not yet demonstrate a substantial final downstream-quality improvement.

Research Question 3 — Does the result generalize beyond 1B?

Hypothesis: The combined architecture remains stable and beneficial at larger model scales.

Results & Analysis:

No completed larger-scale matched comparison is available yet.

Findings: Generalization beyond 1B remains unresolved. The 3B experiment will complete before determining whether the faster early convergence and task-level changes observed at 1B persist at a larger scale. The suspended 7B plan is not currently contributing experimental evidence.

Archive

• WandB Report: https://wandb.ai/archspace/SiameseNormDepthAttention/reports/Siamese-Norm-and-Depth-Attention-in-OLMo-3-1B--VmlldzoxNzYwMzAwNw
• Hugging Face Collection: https://huggingface.co/ArchSpace-Collection/SiameseNorm-DepthAttention

Reviewer Assessment (for repo reviewers)

Draft status remains recommended.

The implementation and 1B end-to-end pipeline are available for review. The matched 1B experiment indicates faster early convergence, but the final eight-task macro average is nearly tied with the OLMo 3 baseline.

The modified 3B experiment remains in progress, while the 7B experimental plan is suspended. Larger-scale validation is therefore incomplete, and no general architectural-effectiveness claim is made at this stage.

Merge Checklist:

• The implementation PR is linked to an in-progress Architecture Proposal issue ([ARCH-PROP] SiameseNorm&Depth-Attention #2).
• Implementation correctness is verified at 1B through four-stage training, checkpoint-transition, and inference validation.
• A matched OLMo 3 1B baseline comparison has been completed.
• The modified 3B training and downstream evaluation are complete.
• The effectiveness of the proposed architecture is validated at the agreed larger-model scale.
• The archive information is complete, including final WandB reports and Hugging Face collections for all required model scales.